import cv2
import numpy as np
import pyautogui
import pydirectinput
from ultralytics import YOLO
import time
import heapq

# A*寻路算法的节点类
class Node:
    def __init__(self, x, y, parent=None):
        self.x = x
        self.y = y
        self.parent = parent
        self.g = 0  # 从起点到当前节点的实际代价
        self.h = 0  # 从当前节点到目标节点的估计代价
        self.f = 0  # 总代价 f = g + h

    def __eq__(self, other):
        return self.x == other.x and self.y == other.y

    def __lt__(self, other):
        return self.f < other.f

    def __hash__(self):
        return hash((self.x, self.y))

# 改进的A*寻路算法实现，支持对角线移动
def astar_search(grid, start, end):
    # 检查起点和终点是否在网格范围内且可达
    if not (0 <= start[0] < len(grid) and 0 <= start[1] < len(grid[0]) and
            0 <= end[0] < len(grid) and 0 <= end[1] < len(grid[0])):
        return None
    
    if grid[start[0]][start[1]] or grid[end[0]][end[1]]:
        return None
    
    # 定义相邻节点的偏移量（包括对角线）
    neighbors = [(-1, 0), (1, 0), (0, -1), (0, 1), (-1, -1), (-1, 1), (1, -1), (1, 1)]
    
    # 初始化起点和终点节点
    start_node = Node(start[0], start[1])
    end_node = Node(end[0], end[1])
    
    # 开放列表和关闭列表
    open_list = []
    closed_set = set()
    
    # 将起点加入开放列表
    heapq.heappush(open_list, (start_node.f, start_node))
    
    while open_list:
        # 获取f值最小的节点
        _, current_node = heapq.heappop(open_list)
        
        # 如果到达终点，回溯路径
        if current_node == end_node:
            path = []
            while current_node:
                path.append((current_node.x, current_node.y))
                current_node = current_node.parent
            return path[::-1]  # 反转路径
        
        # 将当前节点加入关闭列表
        closed_set.add((current_node.x, current_node.y))
        
        # 检查所有相邻节点
        for dx, dy in neighbors:
            neighbor_x = current_node.x + dx
            neighbor_y = current_node.y + dy
            
            # 检查是否超出边界或不可通行
            if not (0 <= neighbor_x < len(grid) and 0 <= neighbor_y < len(grid[0])):
                continue
            
            # 检查对角线移动时的两个相邻格子是否都可通行
            if dx != 0 and dy != 0:
                if grid[current_node.x][neighbor_y] or grid[neighbor_x][current_node.y]:
                    continue
            
            if grid[neighbor_x][neighbor_y] or (neighbor_x, neighbor_y) in closed_set:
                continue
            
            # 创建相邻节点
            neighbor_node = Node(neighbor_x, neighbor_y, current_node)
            
            # 计算g值（对角线移动代价更高）
            if dx != 0 and dy != 0:
                neighbor_node.g = current_node.g + 1.4  # 对角线移动代价
            else:
                neighbor_node.g = current_node.g + 1  # 直线移动代价
            
            # 计算h值（使用对角线距离，更精确）
            dx_h = abs(neighbor_x - end_node.x)
            dy_h = abs(neighbor_y - end_node.y)
            neighbor_node.h = (dx_h + dy_h) + (np.sqrt(2) - 2) * min(dx_h, dy_h)
            
            # 计算f值
            neighbor_node.f = neighbor_node.g + neighbor_node.h
            
            # 检查是否已在开放列表中且有更好的路径
            in_open = False
            for _, node in open_list:
                if node == neighbor_node and node.g <= neighbor_node.g:
                    in_open = True
                    break
            
            if not in_open:
                heapq.heappush(open_list, (neighbor_node.f, neighbor_node))
    
    # 没有找到路径
    return None

# 路径平滑处理
def smooth_path(path, grid):
    if len(path) < 3:
        return path
    
    smoothed_path = [path[0]]
    current_idx = 0
    
    while current_idx < len(path) - 1:
        for end_idx in range(len(path) - 1, current_idx, -1):
            if is_line_passable(path[current_idx], path[end_idx], grid):
                smoothed_path.append(path[end_idx])
                current_idx = end_idx
                break
    
    return smoothed_path

# 检查两点之间的直线是否可通行
def is_line_passable(p1, p2, grid):
    x1, y1 = p1
    x2, y2 = p2
    
    dx = abs(x2 - x1)
    dy = abs(y2 - y1)
    sx = 1 if x1 < x2 else -1
    sy = 1 if y1 < y2 else -1
    err = dx - dy
    
    while True:
        if grid[x1][y1]:
            return False
        
        if x1 == x2 and y1 == y2:
            break
            
        e2 = 2 * err
        if e2 > -dy:
            err -= dy
            x1 += sx
        if e2 < dx:
            err += dx
            y1 += sy
    
    return True

# 创建游戏地图网格
def create_grid(screen_width, screen_height, cell_size, obstacles):
    # 计算网格尺寸
    grid_width = screen_width // cell_size
    grid_height = screen_height // cell_size
    
    # 初始化全为False的网格（False表示可通行）
    grid = [[False for _ in range(grid_width)] for _ in range(grid_height)]
    
    # 将障碍物区域标记为True（不可通行）
    for obstacle in obstacles:
        x1, y1, x2, y2 = obstacle
        
        # 计算障碍物在网格中的位置
        start_x = max(0, x1 // cell_size)
        end_x = min(grid_width - 1, x2 // cell_size)
        start_y = max(0, y1 // cell_size)
        end_y = min(grid_height - 1, y2 // cell_size)
        
        # 将障碍物区域标记为不可通行
        for y in range(start_y, end_y + 1):
            for x in range(start_x, end_x + 1):
                grid[y][x] = True
    
    return grid

# 将世界坐标转换为网格坐标
def world_to_grid(world_pos, cell_size):
    return (world_pos[1] // cell_size, world_pos[0] // cell_size)

# 将网格坐标转换为世界坐标（返回网格中心）
def grid_to_world(grid_pos, cell_size):
    return (grid_pos[1] * cell_size + cell_size // 2, 
            grid_pos[0] * cell_size + cell_size // 2)

# 根据路径执行移动，增加方向稳定性检查
def follow_path(path, cell_size, current_pos, last_direction, direction_counter):
    if not path or len(path) <= 1:
        return last_direction, direction_counter
    
    # 获取下一个目标网格位置
    next_grid_pos = path[1]  # path[0]是当前位置
    
    # 转换为世界坐标
    next_world_pos = grid_to_world(next_grid_pos, cell_size)
    
    # 计算当前位置和目标位置的中心
    current_center_x = (current_pos[0] + current_pos[2]) // 2
    current_center_y = (current_pos[1] + current_pos[3]) // 2
    
    # 计算移动方向
    dx = next_world_pos[0] - current_center_x
    dy = next_world_pos[1] - current_center_y
    
    # 根据方向优先级选择移动按键
    new_direction = None
    if abs(dx) > abs(dy):
        # 水平方向优先
        if dx > 10:  # 设置一个阈值，避免微小移动
            new_direction = 'right'
        elif dx < -10:
            new_direction = 'left'
        elif dy > 10:
            new_direction = 'down'
        elif dy < -10:
            new_direction = 'up'
    else:
        # 垂直方向优先
        if dy > 10:
            new_direction = 'down'
        elif dy < -10:
            new_direction = 'up'
        elif dx > 10:
            new_direction = 'right'
        elif dx < -10:
            new_direction = 'left'
    
    # 方向稳定性检查：避免短时间内频繁改变方向
    if new_direction and new_direction != last_direction:
        direction_counter += 1
        # 如果频繁改变方向，保持原方向
        if direction_counter >= 2:
            new_direction = last_direction
            direction_counter = 0
    else:
        direction_counter = 0
    
    # 执行移动
    if new_direction:
        pydirectinput.press(new_direction)
        time.sleep(0.1)
    
    return new_direction, direction_counter

# 可视化网格和路径（调试用）
def visualize_grid_and_path(grid, path, self_pos, ore_pos, cell_size, exploration_target=None):
    # 创建空白图像
    img = np.zeros((len(grid) * cell_size, len(grid[0]) * cell_size, 3), dtype=np.uint8)
    
    # 绘制障碍物
    for y in range(len(grid)):
        for x in range(len(grid[0])):
            if grid[y][x]:
                cv2.rectangle(img, (x*cell_size, y*cell_size), 
                             ((x+1)*cell_size, (y+1)*cell_size), 
                             (0, 0, 255), -1)  # 红色表示障碍物
    
    # 绘制路径
    if path:
        for i in range(len(path) - 1):
            start = (path[i][1] * cell_size + cell_size//2, 
                    path[i][0] * cell_size + cell_size//2)
            end = (path[i+1][1] * cell_size + cell_size//2, 
                  path[i+1][0] * cell_size + cell_size//2)
            cv2.line(img, start, end, (0, 255, 0), 2)  # 绿色线条表示路径
    
    # 绘制自身位置
    if self_pos:
        center_x = (self_pos[0] + self_pos[2]) // 2 // cell_size
        center_y = (self_pos[1] + self_pos[3]) // 2 // cell_size
        cv2.circle(img, (center_x*cell_size + cell_size//2, 
                        center_y*cell_size + cell_size//2), 
                  cell_size//3, (255, 0, 0), -1)  # 蓝色圆圈表示自身位置
    
    # 绘制矿石位置
    if ore_pos:
        center_x = (ore_pos[0] + ore_pos[2]) // 2 // cell_size
        center_y = (ore_pos[1] + ore_pos[3]) // 2 // cell_size
        cv2.circle(img, (center_x*cell_size + cell_size//2, 
                        center_y*cell_size + cell_size//2), 
                  cell_size//3, (255, 255, 0), -1)  # 青色圆圈表示矿石位置
    
    # 绘制探索目标
    if exploration_target:
        cv2.circle(img, (exploration_target[1]*cell_size + cell_size//2, 
                        exploration_target[0]*cell_size + cell_size//2), 
                  cell_size//3, (0, 165, 255), -1)  # 橙色圆圈表示探索目标
    
    # 调整图像大小以便于查看
    img = cv2.resize(img, (img.shape[1]*2, img.shape[0]*2), interpolation=cv2.INTER_NEAREST)
    
    return img

# 生成探索目标
def generate_exploration_target(grid, current_pos, cell_size, exploration_count, last_exploration_target=None):
    grid_height = len(grid)
    grid_width = len(grid[0])
    
    # 将当前位置转换为网格坐标
    current_grid_pos = world_to_grid(((current_pos[0] + current_pos[2])//2, 
                                     (current_pos[1] + current_pos[3])//2), cell_size)
    
    # 简单的探索策略：按顺时针方向探索地图边缘
    # 根据探索计数确定目标方向
    direction = exploration_count % 4
    
    # 根据方向选择探索目标
    if direction == 0:  # 上方
        target_x = current_grid_pos[1]
        target_y = 1  # 稍微离开顶部边缘
    elif direction == 1:  # 右方
        target_x = grid_width - 2  # 稍微离开右边缘
        target_y = current_grid_pos[0]
    elif direction == 2:  # 下方
        target_x = current_grid_pos[1]
        target_y = grid_height - 2  # 稍微离开底部边缘
    else:  # 左方
        target_x = 1  # 稍微离开左边缘
        target_y = current_grid_pos[0]
    
    # 如果目标位置不可达，尝试寻找附近的可达位置
    if grid[target_y][target_x]:
        for dy in [-1, 0, 1]:
            for dx in [-1, 0, 1]:
                if 0 <= target_y+dy < grid_height and 0 <= target_x+dx < grid_width:
                    if not grid[target_y+dy][target_x+dx]:
                        target_x += dx
                        target_y += dy
                        break
            else:
                continue
            break
    
    return (target_y, target_x)

def capture_screen(x, y, width, height):
    """捕获指定区域的屏幕截图"""
    screenshot = pyautogui.screenshot(region=(x, y, width, height))
    return cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)

def detect_objects(image, model):
    """使用 YOLOv8 模型对图像进行目标检测"""
    return model(image)

def print_detection_results(results):
    """解析检测结果并返回位置列表"""
    detection_info = {
        '矿': [],
        '障碍物': [],
        '自己': []
    }
    for r in results:
        boxes = r.boxes
        for box in boxes:
            class_id = r.names[int(box.cls)]
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            if class_id in detection_info:
                detection_info[class_id].append((x1, y1, x2, y2))
    return detection_info['矿'], detection_info['自己'], detection_info['障碍物']

def calculate_distance(box1, box2):
    """计算两个边界框中心点的距离"""
    x1_center = (box1[0] + box1[2]) / 2
    y1_center = (box1[1] + box1[3]) / 2
    x2_center = (box2[0] + box2[2]) / 2
    y2_center = (box2[1] + box2[3]) / 2
    return np.sqrt((x1_center - x2_center) ** 2 + (y1_center - y2_center) ** 2)

def is_colliding(box1, box2):
    """检查两个边界框是否相交"""
    x1_1, y1_1, x2_1, y2_1 = box1
    x1_2, y1_2, x2_2, y2_2 = box2
    return not (x2_1 < x1_2 or x2_2 < x1_1 or y2_1 < y1_2 or y2_2 < y1_1)

def avoid_lava(self_position, lava_positions):
    """岩浆躲避逻辑"""
    for lava in lava_positions:
        if is_colliding(self_position, lava):
            print("警告：与岩浆碰撞，执行躲避")
            lava_center_x = (lava[0] + lava[2]) / 2
            self_center_x = (self_position[0] + self_position[2]) / 2
            # 优先水平躲避
            pydirectinput.press('left' if self_center_x < lava_center_x else 'right')
            time.sleep(0.1)
            return True
    return False

def move_to_ore_step(ores, self_positions, obstacles, x, y, width, height, model, exploration_count, exploration_target=None, last_direction=None, direction_counter=0):
    """执行单步移动并返回新状态"""
    if not self_positions:
        print("未检测到自身位置")
        return [], [], [], exploration_count, exploration_target, last_direction, direction_counter
    
    self_position = self_positions[0]
    screen_rect = (x, y, x + width, y + height)
    
    # 创建游戏地图网格
    cell_size = 20  # 网格大小，单位为像素
    grid = create_grid(width, height, cell_size, obstacles)
    
    # 如果有矿石，前往最近的矿石
    if ores:
        # 选择最近的矿石
        distances = [calculate_distance(ore, self_position) for ore in ores]
        closest_idx = np.argmin(distances)
        target_ore = ores[closest_idx]
        target_distance = distances[closest_idx]

        # 检查是否已达到挖矿距离
        if target_distance < 100:
            print(f"已接近目标，距离: {target_distance:.1f}px，准备挖矿")
            return ores, self_positions, obstacles, exploration_count, exploration_target, last_direction, direction_counter
        
        # 计算自身和矿石在网格中的位置
        self_grid_pos = world_to_grid(((self_position[0] + self_position[2])//2, 
                                      (self_position[1] + self_position[3])//2), cell_size)
        ore_grid_pos = world_to_grid(((target_ore[0] + target_ore[2])//2, 
                                     (target_ore[1] + target_ore[3])//2), cell_size)
        
        # 使用A*算法寻找路径
        path = astar_search(grid, self_grid_pos, ore_grid_pos)
        
        # 平滑路径
        if path:
            path = smooth_path(path, grid)
        
        # 可视化网格和路径（调试用）
        if path:
            grid_img = visualize_grid_and_path(grid, path, self_position, target_ore, cell_size)
            cv2.imshow('Pathfinding Grid', grid_img)
            cv2.waitKey(1)
        
        # 沿着路径移动
        if path and len(path) > 1:
            print(f"找到路径，长度: {len(path)}，前往矿石")
            last_direction, direction_counter = follow_path(path, cell_size, self_position, last_direction, direction_counter)
        else:
            print("未找到通往矿石的有效路径，尝试简单移动")
            # 如果没有找到路径，使用原有的简单移动逻辑
            directions = get_move_directions(target_ore, self_position)
            print(f"目标距离: {target_distance:.1f}px, 移动方向优先级: {directions}")
            
            new_direction = None
            for direction in directions:
                new_pos = list(self_position)
                step_size = 80  # 默认步长

                # 根据方向调整移动参数
                if direction == 'right':
                    new_pos[0] += step_size
                    new_pos[2] += step_size
                elif direction == 'left':
                    new_pos[0] -= step_size
                    new_pos[2] -= step_size
                elif direction == 'down':
                    new_pos[1] += step_size
                    new_pos[3] += step_size
                elif direction == 'up':
                    new_pos[1] -= step_size
                    new_pos[3] -= step_size

                # 边界检查
                if (new_pos[0] < screen_rect[0] or new_pos[2] > screen_rect[2] or
                        new_pos[1] < screen_rect[1] or new_pos[3] > screen_rect[3]):
                    print(f"方向 {direction} 超出边界，跳过")
                    continue

                # 碰撞检查
                if any(is_colliding(new_pos, obs) for obs in obstacles):
                    print(f"方向 {direction} 被障碍物阻挡，尝试下一个方向")
                    continue

                # 执行移动
                print(f"执行移动: {direction}")
                pydirectinput.press(direction)
                time.sleep(0.1)
                new_direction = direction
                break
            else:
                print("所有方向都无法移动，随机尝试")
                random_dir = np.random.choice(['up', 'down', 'left', 'right'])
                pydirectinput.press(random_dir)
                time.sleep(0.1)
                new_direction = random_dir
            
            # 更新方向状态
            if new_direction and new_direction != last_direction:
                direction_counter += 1
                if direction_counter >= 2:
                    new_direction = last_direction
                    direction_counter = 0
            else:
                direction_counter = 0
            
            last_direction = new_direction
    
    # 如果没有矿石，执行探索
    else:
        print("未检测到矿石，开始探索")
        
        # 如果没有探索目标，生成一个
        if exploration_target is None:
            exploration_target = generate_exploration_target(grid, self_position, cell_size, exploration_count)
            exploration_count += 1
            print(f"生成新的探索目标: {exploration_target}")
        
        # 计算自身和探索目标在网格中的位置
        self_grid_pos = world_to_grid(((self_position[0] + self_position[2])//2, 
                                      (self_position[1] + self_position[3])//2), cell_size)
        
        # 使用A*算法寻找路径
        path = astar_search(grid, self_grid_pos, exploration_target)
        
        # 平滑路径
        if path:
            path = smooth_path(path, grid)
        
        # 可视化网格和路径（调试用）
        if path:
            grid_img = visualize_grid_and_path(grid, path, self_position, None, cell_size, exploration_target)
            cv2.imshow('Pathfinding Grid', grid_img)
            cv2.waitKey(1)
        
        # 沿着路径移动
        if path and len(path) > 1:
            print(f"找到路径，长度: {len(path)}，前往探索目标")
            last_direction, direction_counter = follow_path(path, cell_size, self_position, last_direction, direction_counter)
            
            # 检查是否接近探索目标
            distance_to_target = calculate_distance(
                (self_position[0], self_position[1], self_position[2], self_position[3]),
                (exploration_target[1]*cell_size, exploration_target[0]*cell_size, 
                 exploration_target[1]*cell_size+cell_size, exploration_target[0]*cell_size+cell_size)
            )
            
            # 使用hysteresis机制稳定目标检测
            if distance_to_target < 80:  # 降低阈值，避免频繁触发
                print("到达探索目标附近，准备生成新目标")
                if distance_to_target < 50:  # 更小的阈值确认到达
                    print("确认到达探索目标，生成新目标")
                    exploration_target = generate_exploration_target(grid, self_position, cell_size, exploration_count)
                    exploration_count += 1
        else:
            print("未找到通往探索目标的有效路径，生成新目标")
            exploration_target = generate_exploration_target(grid, self_position, cell_size, exploration_count)
            exploration_count += 1
    
    # 移动后立即检测新状态
    new_screenshot = capture_screen(x, y, width, height)
    new_results = detect_objects(new_screenshot, model)
    new_ores, new_self_positions, new_obstacles = print_detection_results(new_results)
    
    return new_ores, new_self_positions, new_obstacles, exploration_count, exploration_target, last_direction, direction_counter

def get_move_directions(ore, self_position):
    """根据矿石位置计算移动方向优先级"""
    ore_center_x, ore_center_y = (ore[0] + ore[2]) / 2, (ore[1] + ore[3]) / 2
    self_center_x, self_center_y = (self_position[0] + self_position[2]) / 2, (self_position[1] + self_position[3]) / 2

    # 计算水平和垂直距离差
    dx = ore_center_x - self_center_x
    dy = ore_center_y - self_center_y

    # 优先移动距离差较大的方向
    directions = []
    if abs(dx) > abs(dy):
        directions.append('right' if dx > 0 else 'left')
        directions.append('down' if dy > 0 else 'up')
    else:
        directions.append('down' if dy > 0 else 'up')
        directions.append('right' if dx > 0 else 'left')

    # 添加次要方向
    directions.append('left' if dx < 0 else 'right')
    directions.append('up' if dy < 0 else 'down')

    return directions

def main():
    x, y, width, height = 0, 87, 780, 660
    model = YOLO('../runs/best.pt')
    
    # 初始化探索计数器和目标
    exploration_count = 0
    exploration_target = None
    
    # 初始化方向状态
    last_direction = None
    direction_counter = 0

    print("开始自动寻矿...")
    print("按 Q 键退出程序")

    try:
        while True:
            # 捕获屏幕并检测
            screenshot = capture_screen(x, y, width, height)
            results = detect_objects(screenshot, model)
            ores, self_positions, obstacles = print_detection_results(results)

            # 检查是否需要躲避岩浆
            if self_positions and obstacles:
                if avoid_lava(self_positions[0], obstacles):
                    continue  # 躲避后重新检测

            # 处理矿石或探索
            if ores or self_positions:
                # 执行单步移动并更新状态
                ores, self_positions, obstacles, exploration_count, exploration_target, last_direction, direction_counter = move_to_ore_step(
                    ores, self_positions, obstacles, x, y, width, height, model, exploration_count, exploration_target, last_direction, direction_counter
                )

                # 检查是否到达挖矿距离
                if ores and self_positions:
                    distance = calculate_distance(ores[0], self_positions[0])
                    if distance < 100:
                        print(f"到达挖矿距离，距离: {distance:.1f}px")
                        ore_center_x = int((ores[0][0] + ores[0][2]) / 2)
                        ore_center_y = int((ores[0][1] + ores[0][3]) / 2 + 80)

                        # 点击采矿
                        print(f"开始采矿，坐标: ({ore_center_x}, {ore_center_y})")
                        for _ in range(7):
                            pydirectinput.click(ore_center_x, ore_center_y)
                            

            # 显示当前屏幕
            cv2.imshow('Game Screen', screenshot)

            # 退出条件
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    except KeyboardInterrupt:
        print("\n程序被用户中断")
    finally:
        cv2.destroyAllWindows()
        print("程序已退出")

if __name__ == "__main__":
    main()